// Copyright 2023 Srhino Co., Ltd. All rights reserved.

package core

import (
	"context"
	"git.ouryun.cn/lzj1/grpc-bidirectional-stream/pkg/apis/v3"
	"git.ouryun.cn/lzj1/grpc-bidirectional-stream/pkg/utils"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"net"
	"sync"
)

type grpcServer struct {
	v3.UnimplementedCoreServiceServer
	// {"引擎标识":{"组件标识":stream}}
	streams map[string]map[string]v3.CoreService_AgentServer
	ctx     context.Context
	config  Config
	mu      sync.RWMutex
}

func (g *grpcServer) Agent(stream v3.CoreService_AgentServer) error {
	defer func() {
		if r := recover(); r != nil {
			utils.Logger.Error(r)
		}
	}()
	utils.Logger.Info("client connection is successful")
	// 获取客户端的唯一标识符，可以根据实际需求进行处理
	md, ok := metadata.FromIncomingContext(stream.Context())
	if !ok {
		return status.Errorf(codes.InvalidArgument, "missing metadata")
	}
	// 获取引擎标识
	e := md.Get(ENGINE_KEY)
	if len(e) == 0 {
		return status.Errorf(codes.InvalidArgument, "missing %s in metadata", ENGINE_KEY)
	}
	engine := e[0]
	utils.Logger.Infof("engine name id is %s", engine)
	// 获取组件标识
	c := md.Get(COMPONENT_KEY)
	if len(c) == 0 {
		return status.Errorf(codes.InvalidArgument, "missing %s in metadata", COMPONENT_KEY)
	}
	componentName := c[0]
	utils.Logger.Infof("component name id is %s", componentName)

	g.streams[clientId] = stream

	// 接收客户端数据
	err := s.doRecv(stream)
	if err != nil {
		return err
	}
	return nil
}

func (g *grpcServer) Start() error {
	listener, _ := net.Listen("tcp", g.config.GetAddress())
	grpcserver := grpc.NewServer()
	// 注册服务方法
	v3.RegisterCoreServiceServer(grpcserver, g)

	if err := grpcserver.Serve(listener); err != nil {
		utils.Logger.Info(err)
		return err
	}

	return nil
}

func (g *grpcServer) Stop() {
	//TODO implement me
	panic("implement me")
}

func NewGrpcServer(ctx context.Context, config Config) GrpcServer {
	return &grpcServer{
		streams: make(map[string]map[string]v3.CoreService_AgentServer),
		ctx:     ctx,
		config:  config,
	}
}

func (g *grpcServer) addStream(engine string, componentName string, stream v3.CoreService_AgentServer) error {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.streams[engine]; !ok {
		g.streams[engine] = make(map[string]v3.CoreService_AgentServer)
	}
	g.streams[engine][componentName] = stream
	return nil
}

func (g *grpcServer) getStream(engine string, componentName string) v3.CoreService_AgentServer {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.streams[engine]; !ok {
		g.streams[engine] = make(map[string]v3.CoreService_AgentServer)
	}
	g.streams[engine][componentName] = stream
	return nil
}
